!> author: 左志华
!> date: 2022/5/11
!>
!> Statistics <br>
!> 统计学函数
module easy_stats_m

    use, intrinsic :: iso_fortran_env, only: rk => real64
    implicit none
    private

    public :: mean, var, cov, randu, randn

    !> 平均数
    interface mean
        procedure :: mean_real, mean_integer
    end interface mean

    !> 方差
    interface var
        procedure :: var_real, var_integer
    end interface var

contains

    !> Compute the mean of a vector of single precision real numbers
    pure function mean_real(x) result(y)
        real(rk), intent(in) :: x(..)
        real(rk) :: y

        include "mean.finc"
    end function mean_real

    !> Compute the mean of a vector of integer numbers
    pure function mean_integer(x) result(y)
        integer, intent(in) :: x(..)
        integer :: y

        include "mean.finc"
    end function mean_integer

    !> Compute the variance of a vector of single precision real numbers
    pure function var_real(x, normalized) result(y)
        real(rk), intent(in) :: x(..)
        logical, optional, intent(in) :: normalized !! by default, .false.
        real(rk) :: y
        logical :: normalized_

        if (present(normalized)) then
            normalized_ = normalized
        else
            normalized_ = .false.
        end if

        include "var.finc"
    end function var_real

    !> Compute the variance of a vector of integer numbers
    pure function var_integer(x, normalized) result(y)
        integer, intent(in) :: x(..)
        logical, optional, intent(in) :: normalized !! by default, .false.
        integer :: y
        logical :: normalized_

        if (present(normalized)) then
            normalized_ = normalized
        else
            normalized_ = .false.
        end if

        include "var.finc"
    end function var_integer

    !> generates a random number uniformly distributed in [0, 1).
    impure elemental subroutine randu(x)
        real(rk), intent(out) :: x

        call random_number(x)
    end subroutine randu

    !> generates a random number normally distributed with mean 0 and variance 1.
    impure elemental subroutine randn(x)
        real(rk), intent(out) :: x
        real(rk) :: u
        logical, save :: first_run = .true.
        real(rk), save :: v, sln
        real(rk), parameter :: vsmall = tiny(1.0_rk)

        if (first_run) then
            first_run = .false.
            do
                call random_number(u)
                call random_number(v)
                u = 2*u - 1
                v = 2*v - 1
                x = u*u + v*v + vsmall
                if (x < 1) exit
            end do
            sln = sqrt(-2*log(x)/x)
            x = u*sln
        else
            first_run = .true.
            x = v*sln
        end if

    end subroutine randn

    !> Compute the covariance of a vector of single precision real numbers
    pure function cov(x, normalized) result(y)
        real(rk), intent(in) :: x(:)
        logical, optional, intent(in) :: normalized !! by default, .false.
        real(rk) :: y
        logical :: normalized_
        real(rk) :: x_(size(x))

        if (present(normalized)) then
            normalized_ = normalized
        else
            normalized_ = .false.
        end if

        x_ = x - mean(x)
        y = dot_product(x_, x_)/merge(size(x) - 1, size(x), normalized_)
    end function cov

end module easy_stats_m
